# 数据增强效果
# %cd PaddleClas/
from ppcls.data.imaug import DecodeImage
from ppcls.data.imaug import ResizeImage
from ppcls.data.imaug import RandAugment
from ppcls.data.imaug import transform
from ppcls.data.imaug import Cutout
import os
from matplotlib import pyplot as plt
from PIL import Image
import cv2


size = 224
imgs = []
decode_op = DecodeImage()
resize_op = ResizeImage(size=(size, size))
# randaugment_op = RandAugment()
cutout_op = Cutout(n_holes=1, length=112)

# ops = [decode_op, resize_op, randaugment_op]
ops = [decode_op, resize_op, cutout_op]

imgs_dir = "./dataset/dtd/images/smeared/"
save_path_1 = "./dataAugm/dtd"
fnames = os.listdir(imgs_dir)
for f in fnames:
    data = open(os.path.join(imgs_dir, f),"rb").read()
    img = transform(data, ops)
    save_path = os.path.join(save_path_1,f)
    cv2.imwrite(save_path, img)
    imgs.append((img, imgs_dir))

a, ax = plt.subplots(3, 3, figsize=(12,12))
for i, img in enumerate(imgs[:9]):
    ax[i//3, i%3].imshow(img[0])
    ax[i//3, i%3].axis('off')
    ax[i//3, i%3].set_title('label: %s' % img[1])
plt.show()

